/*
 * @Author       : Zou Kai
 * @Date         : 2024-11-18 11:12:37
 * @Last Modified by: Your Name
 * @Last Modified time: 
 * @Description  : 
 */
#include "common.h"
#include "square_projection.h"

SquareProject::SquareProject(){};

bool SquareProject::Init(const std::string extrinic_file, const std::string intrinsic_file, 
            const std::string device_name, const int project_model, const float error_threshold){
    if(error_threshold ==0 ){
        std::cout<<"Error threshold is invalid"<<std::endl;
        return false;
    }
    if(project_model != 0 && project_model != 1){
        std::cout<<"Project model is invalid"<<std::endl;
        return false;
    }
    if(!LoadIntrinsic(intrinsic_file)){
        std::cout<<"Load intrinsic failed"<<std::endl;
        return false;
    };
    if(!LoadExtrinsic(extrinic_file, device_name)){
        std::cout<<"Load extrinsic failed"<<std::endl;
        return false;
    };
    ProjectConfig_.cerr_loss_ = error_threshold;
    ProjectConfig_.project_model_ = project_model;
    return true;
}

bool SquareProject::LoadIntrinsic(const std::string filename){
    if ((access(filename.c_str(), 0)) == -1) {
        std::cout << "Failed to open settings file at: " << filename<<std::endl;;
        return false;
    }                
    YAML::Node cam_config = YAML::LoadFile(filename);
    std::vector<std::vector<double>> K_vec; 
    std::vector<double> D_vec;
    // std::vector<double> resolution;
    D_vec =
        cam_config["intrinsic"]["D"].as<std::vector<double>>();
    K_vec =
        cam_config["intrinsic"]["K"].as<std::vector<std::vector<double>>>();
    ProjectConfig_.resolution_ = cam_config["resolution"].as<std::vector<int>>();
    
    if(!vec2D2Mat(K_vec, ProjectConfig_.camera_matrix_) || !vec1D2Mat(D_vec, ProjectConfig_.distortion_matrix_,0)){
        std::cerr<<"Failed to convert K_vec to camera_matrix"<<std::endl;
        return -1;
    }
    cv::cv2eigen(ProjectConfig_.camera_matrix_, ProjectConfig_.K_matrix_);
    return true;
}


bool SquareProject::LoadExtrinsic(const std::string filename, const std::string camera_name){
    // std::cout<<filename<<std::endl;                    
    YAML::Node lidar_config = YAML::LoadFile(filename);
    if (!lidar_config) {
        std::cout << "Failed to open settings file at: " << filename << std::endl;
        return false;
    }
    Eigen::Quaterniond rotation;
    Eigen::Vector3d t_vec;

    YAML::Node lidarExtr_node = lidar_config["extrinsic"];
    for (auto extr_it = lidarExtr_node.begin(); extr_it != lidarExtr_node.end();
        extr_it++) 
    {
        std::string device_name = (extr_it->first).as<std::string>();
        if(device_name != camera_name){
            continue;
        }
        YAML::Node extr_params = extr_it->second;
        t_vec[0] = extr_params["translation"]["x"].as<double>();
        t_vec[1] = extr_params["translation"]["y"].as<double>();
        t_vec[2] = extr_params["translation"]["z"].as<double>();
        rotation.x() = extr_params["rotation"]["x"].as<double>();
        rotation.y() = extr_params["rotation"]["y"].as<double>();
        rotation.z() = extr_params["rotation"]["z"].as<double>();
        rotation.w() = extr_params["rotation"]["w"].as<double>();
        ProjectConfig_.T_cam_to_lidar_ = toTransfromMatrix(rotation, t_vec);
    }
    return true;
};

bool SquareProject::FarthestPointPixel(pcl::PointCloud<pcl::PointXYZI>::Ptr cloud){
    ProjectConfig_.pixel_grid.assign(ProjectConfig_.resolution_[0]*ProjectConfig_.resolution_[1], {1000,1000,1000});
    for(auto &point:cloud->points){
        pcl::PointXYZI p = point;
        if(std::isnan(p.z) || std::isnan(p.x) || std::isnan(p.y) || (p.z==0)){
            continue;
        }
        Eigen::Vector3d center_point(p.x, p.y, p.z);
        Eigen::Vector3d project_point = ProjectConfig_.K_matrix_ * center_point;
        cv::Point2d pixel(cvRound(project_point[0]/project_point[2]),cvRound( project_point[1]/project_point[2]));
        if(pixel.x<0 || pixel.y<0 || pixel.x>ProjectConfig_.resolution_[0] || pixel.y>ProjectConfig_.resolution_[1]){
            continue;
        }
        int grid;
        if(ProjectConfig_.project_model_==1){ //投影到原图需要加畸变
            std::vector<cv::Point2d> pt_vec;
            pt_vec.push_back(pixel);
            std::vector<cv::Point2d> undis_pt_mat;
            cv::Mat rvec(3, 1, CV_64FC1);
            cv::Mat tvec(3, 1, CV_64FC1);
            rvec = (cv::Mat_<float>(1,3)<<0,0,0);
            tvec = (cv::Mat_<float>(1,3)<<0,0,0);
            cv::fisheye::projectPoints(pt_vec,undis_pt_mat,rvec,tvec,ProjectConfig_.camera_matrix_,ProjectConfig_.distortion_matrix_);
            grid = static_cast<int>(undis_pt_mat[0].y)*ProjectConfig_.resolution_[1] + static_cast<int>(undis_pt_mat[0].x);
        }else{
            grid = static_cast<int>(pixel.y)*ProjectConfig_.resolution_[1] + static_cast<int>(pixel.x);
        }

        if(grid<0 || grid>ProjectConfig_.resolution_[0]*ProjectConfig_.resolution_[1]){
            continue;
        }
        if(ProjectConfig_.pixel_grid[grid][2] > center_point(2) && center_point(2)>0){
            ProjectConfig_.pixel_grid[grid] = center_point;
        }
    }
    return true;
}

std::map<int, std::vector<cv::Point2d>> SquareProject::ErrorProject(){
    std::vector<double> loss_vec({ProjectConfig_.cerr_loss_, -ProjectConfig_.cerr_loss_});
    // std::cout<<loss_vec[0]<<" "<<loss_vec[1]<<std::endl;
    std::map<int,std::vector<cv::Point2d>> point_cerr_loss;
    for(int i = 0; i < ProjectConfig_.pixel_grid.size(); i++){
        Eigen::Vector3d center_point = ProjectConfig_.pixel_grid[i];
        if(center_point[2] ==1000){
            continue;
        }
        // std::cout<<center_point[0]<<" "<<center_point[1]<<" "<<center_point[2]<<std::endl;
        std::vector<cv::Point2d> square_vec;
        for(auto &loss_x:loss_vec){
            double new_x = center_point[0] + loss_x;
            for(auto &loss_y:loss_vec){
                double new_y = center_point[1] + loss_y;
                for(auto &loss_z:loss_vec){
                    double new_z = center_point[2] + loss_z;
                    Eigen::Vector3d new_point(new_x, new_y, new_z);
                    Eigen::Vector3d project_point = ProjectConfig_.K_matrix_ * new_point;
                    cv::Point2d pixel(cvRound(project_point[0]/project_point[2]), cvRound(project_point[1]/project_point[2]));
                    if(ProjectConfig_.project_model_==1){ //投影到原图需要加畸变
                        std::vector<cv::Point2d> pt_vec;
                        pt_vec.push_back(pixel);
                        std::vector<cv::Point2d> undis_pt_mat;
                        cv::Mat rvec(3, 1, CV_64FC1);
                        cv::Mat tvec(3, 1, CV_64FC1);
                        rvec = (cv::Mat_<float>(1,3)<<0,0,0);
                        tvec = (cv::Mat_<float>(1,3)<<0,0,0);
                        cv::fisheye::projectPoints(pt_vec,undis_pt_mat,rvec,tvec,ProjectConfig_.camera_matrix_,ProjectConfig_.distortion_matrix_);
                        square_vec.push_back(undis_pt_mat[0]);
                    }else{
                        square_vec.push_back(pixel);
                    }
                }
            }
        }
        // std::cout<<square_vec<<std::endl;
        point_cerr_loss.insert(std::make_pair(static_cast<int>(i), square_vec));
    }
    return point_cerr_loss;
}